/**
 * @file hs_cuda_pack.cuh
 * @author iam002
 * @brief CUDA接口数据打包, 实现向量化数据访问以提升带宽, 详见 https://zhuanlan.zhihu.com/p/591058808
 * @version 1.0
 * @date 2024-10-15
 * 
 * @copyright Copyright (c) 2024
 * 
 */
#ifndef __HS_CUDA_PACK_CUH__
#define __HS_CUDA_PACK_CUH__

#include <type_traits>
#include <cstdint>

namespace hs
{

namespace cuda
{

/**
 * @brief 每个包支持的最大字节数 16 bytes
 * CUDA 最多支持 128 个 bit 的访问粒度, 1 byte = 8 bit
 */
constexpr int KERNEL_PACK_MAX_BYTES = 128 / 8;

/**
 * @brief 每个包最大数据个数
 * 支持最小数据为 16 bit, 因此 128/16 = 8
 */
constexpr int KERNEL_PACK_MAX_SIZE = 8;

/* 在编译阶段判断两个int常量表达式的最小值 */
constexpr int MIN_CS(int a, int b) 
{
    return a < b ? a : b;
}

/* 为 aligned_storage<>::type 重命名 */
template<int pack_size, typename T>
struct GetPackType
{
    using type = typename std::aligned_storage<pack_size*sizeof(T), pack_size*sizeof(T)>::type;
};

/* 为 aligned_storage<>::type 重命名 */
template<int pack_size, typename T>
using PackType = typename GetPackType<pack_size, T>::type;

/* 用于在编译阶段验证是否内存对齐, 此联合体不会实际使用 */
template<int pack_size, typename T>
union Pack
{   
    /* NOTE MSVC需要预定义宏 _DISABLE_EXTENDED_ALIGNED_STORAGE, 否则编译报错 */
    static_assert(sizeof(PackType<pack_size, T>) == sizeof(T)*pack_size, "");

    __device__ Pack() { }

    PackType<pack_size, T> storage;
    T elem[pack_size];  
};

/* 实际使用的打包类, 其中包含 pack_size 个类型为 T 的数据 */
template<int pack_size, typename T>
struct alignas(sizeof(T) * pack_size) Packed
{
    __device__ Packed() { }

    union {
        T elem[pack_size];
    };
};

/* 获取指定数据类型对应包的元素个数, 返回其中的最小值 */
template<typename T>
constexpr int PackSize()
{
    return MIN_CS( KERNEL_PACK_MAX_BYTES / sizeof(T), KERNEL_PACK_MAX_SIZE );
}
template<typename T, typename U, typename... Args>
constexpr int PackSize()
{
    return MIN_CS( PackSize<T>(), PackSize<U, Args...>() );
}

/* 判断指定数据类型对应的包是否内存对齐 */
template<size_t pack_size>
bool isAlignedForPack()
{
    return true;
}
template<size_t pack_size, typename T, typename... Args>
bool isAlignedForPack(const T* ptr, const Args*... others)
{
    // BUG 不用满足此条件貌似也能运行 ?
    return reinterpret_cast<uintptr_t>(ptr) % sizeof(Pack<pack_size, T>) == 0
        && isAlignedForPack<pack_size, Args...>(others...);
}

} // end namespace

} // end namespace

#endif
